from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict

# 图的整体状态（这是在节点之间共享的公共状态）
class OverallState(TypedDict):
    a: str

# 从 node_1 输出的结果包含不属于整体状态的私有数据
class Node1Output(TypedDict):
    private_data: str

# 私有数据仅在 node_1 和 node_2 之间共享
def node_1(state: OverallState) -> Node1Output:
    output = {"private_data": "set by node_1"}
    print(f"Entered node `node_1`:\n\tInput: {state}.\n\tReturned: {output}")
    return output

# node_2 的输入仅请求在 node_1 之后可用的私有数据
class Node2Input(TypedDict):
    private_data: str

def node_2(state: Node2Input) -> OverallState:
    output = {"a": "set by node_2"}
    print(f"Entered node `node_2`:\n\tInput: {state}.\n\tReturned: {output}")
    return output

# node_3 仅访问整体状态（不访问来自 node_1 的私有数据）
def node_3(state: OverallState) -> OverallState:
    output = {"a": "set by node_3"}
    print(f"Entered node `node_3`:\n\tInput: {state}.\n\tReturned: {output}")
    return output

# 构建状态图
builder = StateGraph(OverallState)
builder.add_node(node_1)  # node_1 是第一个节点
builder.add_node(node_2)  # node_2 是第二个节点，接受来自 node_1 的私有数据
builder.add_node(node_3)  # node_3 是第三个节点，不查看私有数据
builder.add_edge(START, "node_1")  # 从 node_1 开始图形
builder.add_edge("node_1", "node_2")  # 从 node_1 传递到 node_2
builder.add_edge("node_2", "node_3")  # 从 node_2 传递到 node_3（仅共享整体状态）
builder.add_edge("node_3", END)  # 在 node_3 后结束图形
graph = builder.compile()

# 以初始状态调用图形
response = graph.invoke(
    {
        "a": "set at start",
    }
)

print()
print(f"Output of graph invocation: {response}")